{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "TensorFlow Text On Device Guide",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_as3tyDPAvzM"
      },
      "source": [
        "##### Copyright 2021 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-CoWjX1EBXJX"
      },
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7hQmWrtkBBQB"
      },
      "source": [
        "# Converting TensorFlow Text operators to TensorFlow Lite"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qmGnheU8BPKN"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/text/guide/text_tf_lite\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/text/blob/master/docs/guide/text_tf_lite.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/text/blob/master/docs/guide/text_tf_lite.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/text/docs/guide/text_tf_lite.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hz1hOEHPTF2n"
      },
      "source": [
        "## Overview\n",
        "\n",
        "Machine learning models are frequently deployed using TensorFlow Lite to mobile, embedded, and IoT devices to improve data privacy and lower response times. These models often require support for text processing operations. TensorFlow Text version 2.7 and higher provides improved performance, reduced binary sizes, and operations specifically optimized for use in these environments.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_mdIyFfqTMjc"
      },
      "source": [
        "## Text operators\n",
        "\n",
        "The following TensorFlow Text classes can be used from within a TensorFlow Lite model.\n",
        "\n",
        "* `FastWordpieceTokenizer`\n",
        "* `WhitespaceTokenizer`\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x6NAs1fcUwUn"
      },
      "source": [
        "## Model Example"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8ZalFZQvTJf5",
        "outputId": "9e7ba769-9b2c-401f-d26a-353fcf3ced3f"
      },
      "source": [
        "!pip install -q -U tf-nightly\n",
        "!pip install -q -U tensorflow-text-nightly"
      ],
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uL-I0CyPTXnN"
      },
      "source": [
        "from absl import app\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import tensorflow_text as tf_text\n",
        "\n",
        "from tensorflow.lite.python import interpreter"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qj_bJ-xVTfU1"
      },
      "source": [
        "The following code example shows the conversion process and interpretation in Python using a simple test model. Note that the output of a model cannot be a `tf.RaggedTensor` object when you are using TensorFlow Lite. However, you can return the components of a `tf.RaggedTensor` object or convert it using its `to_tensor` function. See [the RaggedTensor guide](https://www.tensorflow.org/guide/ragged_tensor) for more details."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nqQjBcXqTf_0"
      },
      "source": [
        "class TokenizerModel(tf.keras.Model):\n",
        "\n",
        "  def __init__(self, **kwargs):\n",
        "    super().__init__(**kwargs)\n",
        "    self.tokenizer = tf_text.WhitespaceTokenizer()\n",
        "\n",
        "  @tf.function(input_signature=[\n",
        "      tf.TensorSpec(shape=[None], dtype=tf.string, name='input')\n",
        "  ])\n",
        "  def call(self, input_tensor):\n",
        "    return { 'tokens': self.tokenizer.tokenize(input_tensor).flat_values }"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jsPFI-55TiF_",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "721b9c61-feab-4a17-e09f-c9f4775a1ed9"
      },
      "source": [
        "# Test input data.\n",
        "input_data = np.array(['Some minds are better kept apart'])\n",
        "\n",
        "# Define a Keras model.\n",
        "model = TokenizerModel()\n",
        "\n",
        "# Perform TensorFlow Text inference.\n",
        "tf_result = model(tf.constant(input_data))\n",
        "print('TensorFlow result = ', tf_result['tokens'])"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YKpFsvJGTlPq"
      },
      "source": [
        "\n",
        "## Convert the TensorFlow model to TensorFlow Lite\n",
        "\n",
        "When converting a TensorFlow model with TensorFlow Text operators to TensorFlow Lite, you need to\n",
        "indicate to the `TFLiteConverter` that there are custom operators using the\n",
        "`allow_custom_ops` attribute as in the example below. You can then run the model conversion as you normally would. Review the [TensorFlow Lite converter](https://www.tensorflow.org/lite/convert) documentation for a detailed guide on the basics of model conversion."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6hYWezs1Tndo",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "7af5faf6-96a5-41ae-b876-e3934756cd7f"
      },
      "source": [
        "# Convert to TensorFlow Lite.\n",
        "converter = tf.lite.TFLiteConverter.from_keras_model(model)\n",
        "converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]\n",
        "converter.allow_custom_ops = True\n",
        "tflite_model = converter.convert()"
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cxCdhrHATpSR"
      },
      "source": [
        "## Inference\n",
        "\n",
        "For the TensorFlow Lite interpreter to properly read your model containing TensorFlow Text operators, you must configure it to use these custom operators, and provide registration methods for them. Use `tf_text.tflite_registrar.SELECT_TFTEXT_OPS` to provide the full suite of registration functions for the supported TensorFlow Text operators to `InterpreterWithCustomOps`.\n",
        "\n",
        "Note, that while the example below shows inference in Python, the steps are similar in other languages with some minor API translations, and the necessity to build the `tflite_registrar` into your binary. See [TensorFlow Lite Inference](https://www.tensorflow.org/lite/guide/inference) for more details."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kykFg2pXTriw"
      },
      "source": [
        "# Perform TensorFlow Lite inference.\n",
        "interp = interpreter.InterpreterWithCustomOps(\n",
        "    model_content=tflite_model,\n",
        "    custom_op_registerers=tf_text.tflite_registrar.SELECT_TFTEXT_OPS)\n",
        "interp.get_signature_list()"
      ],
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rNGPpHCCTxVX"
      },
      "source": [
        "Next, the TensorFlow Lite interpreter is invoked with the input, providing a result which matches the TensorFlow result from above."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vmSbfbgJTyKY"
      },
      "source": [
        "tokenize = interp.get_signature_runner('serving_default')\n",
        "output = tokenize(input=input_data)\n",
        "print('TensorFlow Lite result = ', output['tokens'])"
      ],
      "execution_count": 7,
      "outputs": []
    }
  ]
}